{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## import modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torch==2.2 in /opt/conda/lib/python3.10/site-packages (2.2.0)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (4.9.0)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (1.12)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (2.8.5)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (3.1.2)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (2023.12.2)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (8.9.2.26)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.3.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (11.0.2.54)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (10.3.2.106)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (11.4.5.107)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.0.106)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.19.3 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (2.19.3)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.105)\n",
      "Requirement already satisfied: triton==2.2.0 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (2.2.0)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12 in /opt/conda/lib/python3.10/site-packages (from nvidia-cusolver-cu12==11.4.5.107->torch==2.2) (12.3.101)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch==2.2) (2.1.3)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.10/site-packages (from sympy->torch==2.2) (1.3.0)\n",
      "Requirement already satisfied: torchvision==0.17.0 in /opt/conda/lib/python3.10/site-packages (0.17.0)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.10/site-packages (from torchvision==0.17.0) (1.26.2)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from torchvision==0.17.0) (2.28.1)\n",
      "Requirement already satisfied: torch==2.2.0 in /opt/conda/lib/python3.10/site-packages (from torchvision==0.17.0) (2.2.0)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /opt/conda/lib/python3.10/site-packages (from torchvision==0.17.0) (9.3.0)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (4.9.0)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (1.12)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (2.8.5)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (3.1.2)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (2023.12.2)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (8.9.2.26)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (12.1.3.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (11.0.2.54)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (10.3.2.106)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (11.4.5.107)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (12.1.0.106)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.19.3 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (2.19.3)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (12.1.105)\n",
      "Requirement already satisfied: triton==2.2.0 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17.0) (2.2.0)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12 in /opt/conda/lib/python3.10/site-packages (from nvidia-cusolver-cu12==11.4.5.107->torch==2.2.0->torchvision==0.17.0) (12.3.101)\n",
      "Requirement already satisfied: charset-normalizer<3,>=2 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision==0.17.0) (2.1.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision==0.17.0) (3.3)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision==0.17.0) (1.26.10)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision==0.17.0) (2022.6.15)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch==2.2.0->torchvision==0.17.0) (2.1.3)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.10/site-packages (from sympy->torch==2.2.0->torchvision==0.17.0) (1.3.0)\n",
      "Requirement already satisfied: matplotlib==3.5.2 in /opt/conda/lib/python3.10/site-packages (3.5.2)\n",
      "Requirement already satisfied: cycler>=0.10 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (4.46.0)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (1.4.5)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (1.26.2)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (21.3)\n",
      "Requirement already satisfied: pillow>=6.2.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (9.3.0)\n",
      "Requirement already satisfied: pyparsing>=2.2.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (3.0.9)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (2.8.2)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib==3.5.2) (1.16.0)\n",
      "Collecting optuna==2.10.1\n",
      "  Downloading optuna-2.10.1-py3-none-any.whl (308 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m308.2/308.2 kB\u001b[0m \u001b[31m12.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: alembic in /opt/conda/lib/python3.10/site-packages (from optuna==2.10.1) (1.13.1)\n",
      "Collecting cliff (from optuna==2.10.1)\n",
      "  Downloading cliff-4.5.0-py3-none-any.whl.metadata (2.0 kB)\n",
      "Collecting cmaes>=0.8.2 (from optuna==2.10.1)\n",
      "  Downloading cmaes-0.10.0-py3-none-any.whl.metadata (19 kB)\n",
      "Requirement already satisfied: colorlog in /opt/conda/lib/python3.10/site-packages (from optuna==2.10.1) (6.8.0)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.10/site-packages (from optuna==2.10.1) (1.26.2)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from optuna==2.10.1) (21.3)\n",
      "Requirement already satisfied: scipy!=1.4.0 in /opt/conda/lib/python3.10/site-packages (from optuna==2.10.1) (1.11.4)\n",
      "Requirement already satisfied: sqlalchemy>=1.1.0 in /opt/conda/lib/python3.10/site-packages (from optuna==2.10.1) (2.0.23)\n",
      "Requirement already satisfied: tqdm in /opt/conda/lib/python3.10/site-packages (from optuna==2.10.1) (4.66.1)\n",
      "Requirement already satisfied: PyYAML in /opt/conda/lib/python3.10/site-packages (from optuna==2.10.1) (6.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging>=20.0->optuna==2.10.1) (3.0.9)\n",
      "Requirement already satisfied: typing-extensions>=4.2.0 in /opt/conda/lib/python3.10/site-packages (from sqlalchemy>=1.1.0->optuna==2.10.1) (4.9.0)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in /opt/conda/lib/python3.10/site-packages (from sqlalchemy>=1.1.0->optuna==2.10.1) (3.0.2)\n",
      "Requirement already satisfied: Mako in /opt/conda/lib/python3.10/site-packages (from alembic->optuna==2.10.1) (1.3.0)\n",
      "Requirement already satisfied: PrettyTable>=0.7.2 in /opt/conda/lib/python3.10/site-packages (from cliff->optuna==2.10.1) (3.9.0)\n",
      "Collecting autopage>=0.4.0 (from cliff->optuna==2.10.1)\n",
      "  Downloading autopage-0.5.2-py3-none-any.whl.metadata (7.9 kB)\n",
      "Collecting cmd2>=1.0.0 (from cliff->optuna==2.10.1)\n",
      "  Downloading cmd2-2.4.3-py3-none-any.whl (147 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m147.2/147.2 kB\u001b[0m \u001b[31m21.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting stevedore>=2.0.1 (from cliff->optuna==2.10.1)\n",
      "  Downloading stevedore-5.1.0-py3-none-any.whl.metadata (2.2 kB)\n",
      "Requirement already satisfied: attrs>=16.3.0 in /opt/conda/lib/python3.10/site-packages (from cmd2>=1.0.0->cliff->optuna==2.10.1) (23.2.0)\n",
      "Collecting pyperclip>=1.6 (from cmd2>=1.0.0->cliff->optuna==2.10.1)\n",
      "  Downloading pyperclip-1.8.2.tar.gz (20 kB)\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: wcwidth>=0.1.7 in /opt/conda/lib/python3.10/site-packages (from cmd2>=1.0.0->cliff->optuna==2.10.1) (0.2.12)\n",
      "Collecting pbr!=2.1.0,>=2.0.0 (from stevedore>=2.0.1->cliff->optuna==2.10.1)\n",
      "  Downloading pbr-6.0.0-py2.py3-none-any.whl.metadata (1.3 kB)\n",
      "Requirement already satisfied: MarkupSafe>=0.9.2 in /opt/conda/lib/python3.10/site-packages (from Mako->alembic->optuna==2.10.1) (2.1.3)\n",
      "Downloading cmaes-0.10.0-py3-none-any.whl (29 kB)\n",
      "Downloading cliff-4.5.0-py3-none-any.whl (81 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m81.1/81.1 kB\u001b[0m \u001b[31m14.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading autopage-0.5.2-py3-none-any.whl (30 kB)\n",
      "Downloading stevedore-5.1.0-py3-none-any.whl (49 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.6/49.6 kB\u001b[0m \u001b[31m8.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading pbr-6.0.0-py2.py3-none-any.whl (107 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m107.5/107.5 kB\u001b[0m \u001b[31m18.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hBuilding wheels for collected packages: pyperclip\n",
      "  Building wheel for pyperclip (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for pyperclip: filename=pyperclip-1.8.2-py3-none-any.whl size=11122 sha256=226e82749e4b03f898079ff5c1356b49c3e1e044ba5edf36b972482d3950b308\n",
      "  Stored in directory: /home/ashish.jha/.cache/pip/wheels/04/24/fe/140a94a7f1036003ede94579e6b4227fe96c840c6f4dcbe307\n",
      "Successfully built pyperclip\n",
      "Installing collected packages: pyperclip, pbr, cmd2, cmaes, autopage, stevedore, cliff, optuna\n",
      "  Attempting uninstall: optuna\n",
      "    Found existing installation: optuna 3.5.0\n",
      "    Uninstalling optuna-3.5.0:\n",
      "      Successfully uninstalled optuna-3.5.0\n",
      "Successfully installed autopage-0.5.2 cliff-4.5.0 cmaes-0.10.0 cmd2-2.4.3 optuna-2.10.1 pbr-6.0.0 pyperclip-1.8.2 stevedore-5.1.0\n"
     ]
    }
   ],
   "source": [
    "!pip install torch==2.2\n",
    "!pip install torchvision==0.17.0\n",
    "!pip install matplotlib==3.5.2\n",
    "!pip install optuna==2.10.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  _torch_pytree._register_pytree_node(\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import optuna\n",
    "\n",
    "torch.use_deterministic_algorithms(True)\n",
    "device = torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## define model architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvNet(nn.Module):\n",
    "    def __init__(self, trial):\n",
    "        super(ConvNet, self).__init__()\n",
    "        num_conv_layers = trial.suggest_int(\"num_conv_layers\", 1, 4)\n",
    "        num_fc_layers = trial.suggest_int(\"num_fc_layers\", 1, 2)\n",
    "\n",
    "        self.layers = []\n",
    "        input_depth = 1 # grayscale image\n",
    "        for i in range(num_conv_layers):\n",
    "            output_depth = trial.suggest_int(f\"conv_depth_{i}\", 16, 64)\n",
    "            self.layers.append(nn.Conv2d(input_depth, output_depth, 3, 1))\n",
    "            self.layers.append(nn.ReLU())\n",
    "            input_depth = output_depth\n",
    "        self.layers.append(nn.MaxPool2d(2))\n",
    "        p = trial.suggest_float(f\"conv_dropout_{i}\", 0.1, 0.4)\n",
    "        self.layers.append(nn.Dropout(p))\n",
    "        self.layers.append(nn.Flatten())\n",
    "\n",
    "        input_feat = self._get_flatten_shape()\n",
    "        for i in range(num_fc_layers):\n",
    "            output_feat = trial.suggest_int(f\"fc_output_feat_{i}\", 16, 64)\n",
    "            self.layers.append(nn.Linear(input_feat, output_feat))\n",
    "            self.layers.append(nn.ReLU())\n",
    "            p = trial.suggest_float(f\"fc_dropout_{i}\", 0.1, 0.4)\n",
    "            self.layers.append(nn.Dropout(p))\n",
    "            input_feat = output_feat\n",
    "        self.layers.append(nn.Linear(input_feat, 10))\n",
    "        self.layers.append(nn.LogSoftmax(dim=1))\n",
    "        \n",
    "        self.model = nn.Sequential(*self.layers)\n",
    "    \n",
    "    def _get_flatten_shape(self):\n",
    "        conv_model = nn.Sequential(*self.layers)\n",
    "        op_feat = conv_model(torch.rand(1, 1, 28, 28))\n",
    "        n_size = op_feat.data.view(1, -1).size(1)\n",
    "        return n_size\n",
    " \n",
    "    def forward(self, x):\n",
    "        return self.model(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## create data loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The mean and standard deviation values are calculated as the mean of all pixel values of all images in the training dataset\n",
    "train_ds = datasets.MNIST('../data', train=True, download=True,\n",
    "                   transform=transforms.Compose([\n",
    "                       transforms.ToTensor(),\n",
    "                       transforms.Normalize((0.1302,), (0.3069,))]))\n",
    "test_ds = datasets.MNIST('../data', train=False, \n",
    "                   transform=transforms.Compose([\n",
    "                       transforms.ToTensor(),\n",
    "                       transforms.Normalize((0.1302,), (0.3069,))]))\n",
    "\n",
    "train_dataloader = torch.utils.data.DataLoader(train_ds, batch_size=32, shuffle=True)\n",
    "test_dataloader = torch.utils.data.DataLoader(test_ds, batch_size=500, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## define training and inference routines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, device, train_dataloader, optim, epoch):\n",
    "    model.train()\n",
    "    for b_i, (X, y) in enumerate(train_dataloader):\n",
    "        X, y = X.to(device), y.to(device)\n",
    "        optim.zero_grad()\n",
    "        pred_prob = model(X)\n",
    "        loss = F.nll_loss(pred_prob, y) # nll is the negative likelihood loss\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "        if b_i % 500 == 0:\n",
    "            print('epoch: {} [{}/{} ({:.0f}%)]\\t training loss: {:.6f}'.format(\n",
    "                epoch, b_i * len(X), len(train_dataloader.dataset),\n",
    "                100. * b_i / len(train_dataloader), loss.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, device, test_dataloader):\n",
    "    model.eval()\n",
    "    loss = 0\n",
    "    success = 0\n",
    "    with torch.no_grad():\n",
    "        for X, y in test_dataloader:\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            pred_prob = model(X)\n",
    "            loss += F.nll_loss(pred_prob, y, reduction='sum').item()  # loss summed across the batch\n",
    "            pred = pred_prob.argmax(dim=1, keepdim=True)  # use argmax to get the most likely prediction\n",
    "            success += pred.eq(y.view_as(pred)).sum().item()\n",
    "\n",
    "    loss /= len(test_dataloader.dataset)\n",
    "    \n",
    "    accuracy = 100. * success / len(test_dataloader.dataset)\n",
    "\n",
    "    print('\\nTest dataset: Overall Loss: {:.4f}, Overall Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
    "        loss, success, len(test_dataloader.dataset), accuracy))\n",
    "    \n",
    "    return accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## define optimizer and model training routine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def objective(trial):\n",
    "    \n",
    "    model = ConvNet(trial)\n",
    "    opt_name = trial.suggest_categorical(\"optimizer\", [\"Adam\", \"Adadelta\", \"RMSprop\", \"SGD\"])\n",
    "    lr = trial.suggest_float(\"lr\", 1e-1, 5e-1, log=True)\n",
    "    optimizer = getattr(optim, opt_name)(model.parameters(), lr=lr)\n",
    "    \n",
    "    for epoch in range(1, 3):\n",
    "        train(model, device, train_dataloader, optimizer, epoch)\n",
    "        accuracy = test(model, device, test_dataloader)\n",
    "        trial.report(accuracy, epoch)\n",
    "        \n",
    "        if trial.should_prune():\n",
    "            raise optuna.exceptions.TrialPruned()\n",
    "\n",
    "    return accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## run the hyperparameter search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[I 2024-02-16 02:15:30,029]\u001b[0m A new study created in memory with name: mastering_pytorch\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 [0/60000 (0%)]\t training loss: 2.279864\n",
      "epoch: 1 [16000/60000 (27%)]\t training loss: 2.347058\n",
      "epoch: 1 [32000/60000 (53%)]\t training loss: 2.261235\n",
      "epoch: 1 [48000/60000 (80%)]\t training loss: 2.291524\n",
      "\n",
      "Test dataset: Overall Loss: 2.3178, Overall Accuracy: 1135/10000 (11%)\n",
      "\n",
      "epoch: 2 [0/60000 (0%)]\t training loss: 2.304916\n",
      "epoch: 2 [16000/60000 (27%)]\t training loss: 2.303748\n",
      "epoch: 2 [32000/60000 (53%)]\t training loss: 2.324766\n",
      "epoch: 2 [48000/60000 (80%)]\t training loss: 2.347335\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[I 2024-02-16 02:16:58,997]\u001b[0m Trial 0 finished with value: 9.58 and parameters: {'num_conv_layers': 2, 'num_fc_layers': 1, 'conv_depth_0': 34, 'conv_depth_1': 41, 'conv_dropout_1': 0.15383558720767948, 'fc_output_feat_0': 62, 'fc_dropout_0': 0.15506861406996786, 'optimizer': 'Adam', 'lr': 0.20023572685997063}. Best is trial 0 with value: 9.58.\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Test dataset: Overall Loss: 2.3296, Overall Accuracy: 958/10000 (10%)\n",
      "\n",
      "epoch: 1 [0/60000 (0%)]\t training loss: 2.297234\n",
      "epoch: 1 [16000/60000 (27%)]\t training loss: 2.321342\n",
      "epoch: 1 [32000/60000 (53%)]\t training loss: 2.248134\n",
      "epoch: 1 [48000/60000 (80%)]\t training loss: 2.282013\n",
      "\n",
      "Test dataset: Overall Loss: 2.4947, Overall Accuracy: 1009/10000 (10%)\n",
      "\n",
      "epoch: 2 [0/60000 (0%)]\t training loss: 2.338793\n",
      "epoch: 2 [16000/60000 (27%)]\t training loss: 2.524946\n",
      "epoch: 2 [32000/60000 (53%)]\t training loss: 2.318933\n",
      "epoch: 2 [48000/60000 (80%)]\t training loss: 2.332036\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[I 2024-02-16 02:19:10,390]\u001b[0m Trial 1 finished with value: 9.8 and parameters: {'num_conv_layers': 3, 'num_fc_layers': 2, 'conv_depth_0': 29, 'conv_depth_1': 49, 'conv_depth_2': 47, 'conv_dropout_2': 0.278335463941041, 'fc_output_feat_0': 51, 'fc_dropout_0': 0.31004103413830064, 'fc_output_feat_1': 58, 'fc_dropout_1': 0.15485667529511574, 'optimizer': 'RMSprop', 'lr': 0.4780748567930119}. Best is trial 1 with value: 9.8.\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Test dataset: Overall Loss: 2.3665, Overall Accuracy: 980/10000 (10%)\n",
      "\n",
      "epoch: 1 [0/60000 (0%)]\t training loss: 2.297451\n",
      "epoch: 1 [16000/60000 (27%)]\t training loss: 2.322544\n",
      "epoch: 1 [32000/60000 (53%)]\t training loss: 2.298779\n",
      "epoch: 1 [48000/60000 (80%)]\t training loss: 2.334569\n",
      "\n",
      "Test dataset: Overall Loss: 2.3327, Overall Accuracy: 958/10000 (10%)\n",
      "\n",
      "epoch: 2 [0/60000 (0%)]\t training loss: 2.291020\n",
      "epoch: 2 [16000/60000 (27%)]\t training loss: 2.341492\n",
      "epoch: 2 [32000/60000 (53%)]\t training loss: 2.317598\n",
      "epoch: 2 [48000/60000 (80%)]\t training loss: 2.369541\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[I 2024-02-16 02:41:10,990]\u001b[0m Trial 2 finished with value: 9.82 and parameters: {'num_conv_layers': 4, 'num_fc_layers': 2, 'conv_depth_0': 34, 'conv_depth_1': 51, 'conv_depth_2': 28, 'conv_depth_3': 55, 'conv_dropout_3': 0.14208996846752817, 'fc_output_feat_0': 35, 'fc_dropout_0': 0.2660624056161214, 'fc_output_feat_1': 36, 'fc_dropout_1': 0.15326719427526703, 'optimizer': 'Adam', 'lr': 0.14119025099148208}. Best is trial 2 with value: 9.82.\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Test dataset: Overall Loss: 2.3158, Overall Accuracy: 982/10000 (10%)\n",
      "\n",
      "epoch: 1 [0/60000 (0%)]\t training loss: 2.295434\n",
      "epoch: 1 [16000/60000 (27%)]\t training loss: 2.309730\n",
      "epoch: 1 [32000/60000 (53%)]\t training loss: 2.298417\n",
      "epoch: 1 [48000/60000 (80%)]\t training loss: 2.383884\n",
      "\n",
      "Test dataset: Overall Loss: 2.3727, Overall Accuracy: 1010/10000 (10%)\n",
      "\n",
      "epoch: 2 [0/60000 (0%)]\t training loss: 2.310924\n",
      "epoch: 2 [16000/60000 (27%)]\t training loss: 2.367006\n",
      "epoch: 2 [32000/60000 (53%)]\t training loss: 2.380785\n",
      "epoch: 2 [48000/60000 (80%)]\t training loss: 2.387081\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[I 2024-02-16 02:45:38,374]\u001b[0m Trial 3 finished with value: 8.92 and parameters: {'num_conv_layers': 3, 'num_fc_layers': 2, 'conv_depth_0': 17, 'conv_depth_1': 29, 'conv_depth_2': 58, 'conv_dropout_2': 0.33864705908649817, 'fc_output_feat_0': 39, 'fc_dropout_0': 0.23750831217409174, 'fc_output_feat_1': 40, 'fc_dropout_1': 0.14041408480339376, 'optimizer': 'Adam', 'lr': 0.310858350845543}. Best is trial 2 with value: 9.82.\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Test dataset: Overall Loss: 2.3431, Overall Accuracy: 892/10000 (9%)\n",
      "\n",
      "epoch: 1 [0/60000 (0%)]\t training loss: 2.330726\n",
      "epoch: 1 [16000/60000 (27%)]\t training loss: 2.365713\n",
      "epoch: 1 [32000/60000 (53%)]\t training loss: 2.304307\n",
      "epoch: 1 [48000/60000 (80%)]\t training loss: 2.256113\n",
      "\n",
      "Test dataset: Overall Loss: 2.3197, Overall Accuracy: 1135/10000 (11%)\n",
      "\n",
      "epoch: 2 [0/60000 (0%)]\t training loss: 2.273054\n",
      "epoch: 2 [16000/60000 (27%)]\t training loss: 2.342382\n",
      "epoch: 2 [32000/60000 (53%)]\t training loss: 2.342956\n",
      "epoch: 2 [48000/60000 (80%)]\t training loss: 2.277427\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[I 2024-02-16 02:49:20,214]\u001b[0m Trial 4 finished with value: 10.32 and parameters: {'num_conv_layers': 2, 'num_fc_layers': 1, 'conv_depth_0': 26, 'conv_depth_1': 35, 'conv_dropout_1': 0.19941309294192455, 'fc_output_feat_0': 59, 'fc_dropout_0': 0.2574163114292745, 'optimizer': 'RMSprop', 'lr': 0.14980606145925923}. Best is trial 4 with value: 10.32.\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Test dataset: Overall Loss: 2.3201, Overall Accuracy: 1032/10000 (10%)\n",
      "\n",
      "results: \n",
      "num_trials_conducted:  5\n",
      "num_trials_pruned:  0\n",
      "num_trials_completed:  5\n",
      "results from best trial:\n",
      "accuracy:  10.32\n",
      "hyperparameters: \n",
      "num_conv_layers: 2\n",
      "num_fc_layers: 1\n",
      "conv_depth_0: 26\n",
      "conv_depth_1: 35\n",
      "conv_dropout_1: 0.19941309294192455\n",
      "fc_output_feat_0: 59\n",
      "fc_dropout_0: 0.2574163114292745\n",
      "optimizer: RMSprop\n",
      "lr: 0.14980606145925923\n"
     ]
    }
   ],
   "source": [
    "study = optuna.create_study(study_name=\"mastering_pytorch\", direction=\"maximize\")\n",
    "study.optimize(objective, n_trials=100, timeout=2000)\n",
    "\n",
    "pruned_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.PRUNED]\n",
    "complete_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]\n",
    "\n",
    "print(\"results: \")\n",
    "print(\"num_trials_conducted: \", len(study.trials))\n",
    "print(\"num_trials_pruned: \", len(pruned_trials))\n",
    "print(\"num_trials_completed: \", len(complete_trials))\n",
    "\n",
    "print(\"results from best trial:\")\n",
    "trial = study.best_trial\n",
    "\n",
    "print(\"accuracy: \", trial.value)\n",
    "print(\"hyperparameters: \")\n",
    "for key, value in trial.params.items():\n",
    "    print(\"{}: {}\".format(key, value))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (Local)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
